'''

langchain tool invoke逻辑
'''
# #
# from langchain_community.tools import BraveSearch
# api_key = "API KEY"
# tool = BraveSearch.from_api_key(api_key=api_key, search_kwargs={"count": 3})
# tool.run("obama middle name")

# from langchain_community.utilities import SearxSearchWrapper
#
# searx = SearxSearchWrapper(searx_host="http://localhost:8888")
# searx.run("What is the capital of France")


# # pip install tavily-python
# import getpass
# import os
#
# os.environ["TAVILY_API_KEY"] = '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'
# if not os.environ.get("TAVILY_API_KEY"):
#     os.environ["TAVILY_API_KEY"] = getpass.getpass("Tavily API key:\n")
from langchain_community.tools import TavilySearchResults
#
# ty_tool = TavilySearchResults(
#     max_results=5,
#     search_depth="advanced",
#     include_answer=True,
#     include_raw_content=True,
#     include_images=True,
#     # include_domains=[...],
#     # exclude_domains=[...],
#     # name="...",            # overwrite default tool name
#     # description="...",     # overwrite default tool description
#     # args_schema=...,       # overwrite default args_schema: BaseModel
# )
#
#
# res = ty_tool.invoke({"query": "上海今天有哪些新闻"})
# # res = ty_tool.run({"query": "上海今天有哪些新闻"})
# print(res)


from langchain_core.tools import  tool


def get_tool_TavilySearchResults():
    # tool1
    from langchain_community.tools.tavily_search import TavilySearchResults
    search = TavilySearchResults(max_results=5)  # pip install langchain-community== 0.3.17
    return search


def get_tool_open_pycharm():
    @tool
    def open_pycharm() -> bool:
        """
        这是一个工具，可以打开pycharm编程软件,工具执行结果为True或false,表示打开成功或打开失败
        """
        # ...
        is_ok = True
        return is_ok
    return open_pycharm

# 算法优化 special
# 训练图片位置；打开训练图片标注查看
def get_tool_open_img():
    @tool
    def open_img(img_path: str) -> bool:
        """
        这是一个工具，用来打开图片,供人类查看核对；输入是图片路径，输出是否打开成功
        """
        # ...
        # 检测环境为linux服务器
        import os
        cmd_str1 = f'display -resize 800 {img_path}' # 宽800
        os.system(cmd_str1)

        # 判断是否成功
        # 策略
        is_ok = True
        return is_ok
    return open_img

# 执行训练脚本 train_seg.sh 根据日志反馈
def get_tool_exe_train():
    @tool
    def exe_train(project_dir: str, data_dir:str) -> bool:
        """
        这是一个工具，用来在训练服务器上执行yolov8模型的训练代码， 输入是项目代码路径，训练数据集路径，输出是否执行训练成功
        """
        # ...
        # 检测环境为linux服务器
        import os
        cmd_str1 = f'cd {project_dir}'
        cmd_str2 = f'nohup python /home/ps/zhangxiancai/yolov8/user_seg_auto_train.py --root_path {data_dir}  >log/v8seg.log 2>&1 &'
        # /home/ps/zhangxiancai/data/241115jzbxh/trainV8Seg_smallBxh
        os.system(cmd_str1)
        os.system(cmd_str2)

        # 判断是否成功
        # 策略 根据日志文件是否有error

        is_ok = True
        return is_ok
    return exe_train


# langchain rag tool
def get_lc_rag_tool():

    import sys
    sys.path.append(r'D:\code\other\LLMs\my_langchain\rag')
    sys.path.append(r'/home/ps/zhangxiancai/llm_deploy/LLMs/my_langchain/rag')
    from rag.rag_lc import get_rag_query
    save_path = f'/home/ps/zhangxiancai/data/llm_deploy/rag_db'
    embd_model_path = r'/mnt/d/xiancai/bigfiles/models/bge-base-zh-v1___5'
    rag_query = get_rag_query(save_path=save_path, embd_model_path=embd_model_path)

    @tool('lc_rag_tool')
    def lc_rag_tool(query:str) -> str:
        """
        这是一个工具，根据查询字符串检索本地数据库，返回相关信息字符串；本地数据库包括各种项目资料和代码
        """
        docs_res = rag_query(query)
        # docs 2 str
        print(len(docs_res))
        str_res = ''
        for ind, d in enumerate(docs_res):
            str_res += f"{ind} {'-' * 50}\n"
            str_res += d.metadata['path'] + '\n'
            str_res += '-' * 10 + '\n'
            str_res += d.page_content + '\n'
        return str_res
    return lc_rag_tool

def test_open_pycharm():
    # print(open_pycharm.name)
    # print(open_pycharm.description)
    # print(open_pycharm.args)
    #
    # res = open_pycharm.invoke({})
    # print(res)

    lrt = get_lc_rag_tool()
    res = lrt.invoke('test fenglun')
    print(res)

if __name__ == '__main__':
    test_open_pycharm()